//
// Copyright 2019 ZetaSQL Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "zetasql/compliance/functions_testlib.h"
#include "zetasql/compliance/functions_testlib_common.h"
#include "zetasql/public/options.pb.h"
#include "zetasql/public/value.h"
#include "zetasql/testing/test_function.h"
#include "zetasql/testing/using_test_value.cc"
#include "absl/strings/substitute.h"

namespace zetasql {
namespace {
constexpr absl::StatusCode OUT_OF_RANGE = absl::StatusCode::kOutOfRange;

// 'sql_standard_mode': if true, uses the SQL2016 standard for JSON function
// names (e.g. JSON_QUERY instead of JSON_EXTRACT).
// 'scalar_test_cases': if true, returns the JSON_VALUE/JSON_EXTRACT_SCALAR test
// cases. Otherwise, returns the JSON_QUERY/JSON_EXTRACT test cases.
// 'json_constructor': function used for constructing json values.
const std::vector<FunctionTestCall> GetJsonTestsCommon(
    bool sql_standard_mode, bool scalar_test_cases,
    const std::function<Value(absl::optional<absl::string_view>)>&
        json_constructor) {
  std::string query_fn_name;
  std::string value_fn_name;
  if (sql_standard_mode) {
    query_fn_name = "json_query";
    value_fn_name = "json_value";
  } else {
    query_fn_name = "json_extract";
    value_fn_name = "json_extract_scalar";
  }

  const Value json1 = json_constructor(
      R"({"a":{"b":[{"c" : "foo", "d": 1.23, "f":null }], "e": true}})");
  const Value json2 = json_constructor(
      R"({"x": [1, 2, 3, 4, 5], "y": [{"a": "bar"}, {"b":"baz"}] })");
  const Value json3 = json_constructor(R"({"a.b.c": 5})");
  const Value json4 = json_constructor(R"({"longer_field_name": []})");
  const Value json5 = json_constructor("true");
  const Value json6 =
      json_constructor(R"({"a":[{"b": [{"c": [{"d": [3]}]}]}]})");
  const Value json7 = json_constructor(R"({"a":{"b": {"c": {"d": 3}}}})");
  const Value json8 = json_constructor(
      R"({"x" : [    ], "y"    :[1,2       ,      5,3  ,4]})");

  // Note: not enclosed in {}.
  const std::string deep_json_string = R"(
  "a" : {
    "b" : {
      "c" : {
        "d" : {
          "e" : {
            "f" : {
              "g" : {
                "h" : {
                  "i" : {
                    "j" : {
                      "k" : {
                        "l" : {
                          "m" : {
                            "x" : "foo",
                            "y" : 10,
                            "z" : [1, 2, 3]
                          }
                        }
                      }
                    }
                  }
                }
              }
            }
          }
        }
      }
    }
  }
  )";
  const std::string wide_json_string = R"(
  "a" : null, "b" : "bar", "c" : false, "d" : [4, 5], "e" : 0.123, "f" : "345",
  "g" : null, "h" : "baz", "i" : true, "j" : [-3, 0], "k" : 0.321, "l" : "678"
  )";

  const Value deep_json =
      json_constructor(absl::StrCat("{", deep_json_string, "}"));
  const Value wide_json =
      json_constructor(absl::StrCat("{", wide_json_string, "}"));

  const int kArrayElements = 20;
  std::vector<int> indexes(kArrayElements);
  std::iota(indexes.begin(), indexes.end(), 0);
  const Value array_of_wide_json = json_constructor(absl::Substitute(
      R"({"arr" : [$0]})",
      absl::StrJoin(
          indexes, ",", [&wide_json_string](std::string* out, int index) {
            absl::StrAppend(out, absl::Substitute(R"({"index" : $0, $1})",
                                                  index, wide_json_string));
          })));
  const Value array_of_deep_json = json_constructor(absl::Substitute(
      R"({"arr" : [$0]})",
      absl::StrJoin(
          indexes, ",", [&deep_json_string](std::string* out, int index) {
            absl::StrAppend(out, absl::Substitute(R"({"index" : $0, $1})",
                                                  index, deep_json_string));
          })));

  std::vector<FunctionTestCall> all_tests;
  if (scalar_test_cases) {
    all_tests = {
        {value_fn_name, {json1, String("$.a.b[0]")}, NullString()},
        {value_fn_name, {json1, String("$.a.b[0].c")}, String("foo")},
        {value_fn_name, {json1, String("$.a.b[0].d")}, String("1.23")},
        {value_fn_name, {json1, String("$.a.b[0].f")}, NullString()},
        {value_fn_name, {json1, String("$.a.e")}, String("true")},
        {value_fn_name, {json2, String("$.y[0].a")}, String("bar")},
        {value_fn_name,
         {json2, absl::StrCat("$.y[0]", EscapeKey(sql_standard_mode, "a"))},
         String("bar")},
        {value_fn_name, {json2, String("$.y[1].b")}, String("baz")},
        {value_fn_name, {json5, String("$")}, String("true")},
        {value_fn_name, {json7, String("$.a.b.c.d")}, String("3")},
        {value_fn_name,
         {json7, absl::StrCat("$", EscapeKey(sql_standard_mode, "a"),
                              EscapeKey(sql_standard_mode, "b"), ".c",
                              EscapeKey(sql_standard_mode, "d"))},
         String("3")},
        {value_fn_name, {json8, String("$.y[3]")}, String("3")},
        // Deep/wide json.
        {value_fn_name,
         {deep_json, String("$.a.b.c.d.e.f.g.h.i.j.k.l.m.x")},
         String("foo")},
        {value_fn_name, {wide_json, String("$.j[1]")}, String("0")},
        {value_fn_name,
         {array_of_deep_json, String("$.arr[13].a.b.c.d.e.f.g.h.i.j.k.l.m.x")},
         String("foo")},
        {value_fn_name,
         {array_of_wide_json, String("$.arr[12].index")},
         String("12")},
        {value_fn_name,
         {array_of_wide_json, String("$.arr[12].h")},
         String("baz")},
        // Non-ASCII UTF-8 and special cases.
        {value_fn_name,
         {json_constructor(R"({"Моша_öá5ホバークラフト鰻鰻" : "x"})"),
          String("$.Моша_öá5ホバークラフト鰻鰻")},
         String("x")},
        {value_fn_name,
         {json_constructor(R"({"1" : 2})"),
          absl::StrCat("$", EscapeKey(sql_standard_mode, "1"))},
         String("2")},
        {value_fn_name,
         {json_constructor(R"({"1" : 2})"), String("$.1")},
         String("2")},
        // Unsupported/unimplemented JSONPath features.
        {value_fn_name,
         {json6, String("$.a[0].b[(@.length-1)].c[(@.length-1)].d")},
         NullString(),
         OUT_OF_RANGE},
    };
    if (sql_standard_mode) {
      all_tests.push_back({value_fn_name,
                           {json3, String("$['a.b.c']")},
                           NullString(),
                           OUT_OF_RANGE});
    } else {
      all_tests.push_back({value_fn_name,
                           {json3, String("$.\"a.b.c\"")},
                           NullString(),
                           OUT_OF_RANGE});
    }
  } else {
    all_tests = {
        {query_fn_name,
         {json_constructor(absl::nullopt), NullString()},
         json_constructor(absl::nullopt)},
        {query_fn_name, {json1, NullString()}, json_constructor(absl::nullopt)},
        {query_fn_name,
         {json_constructor(absl::nullopt), "$"},
         json_constructor(absl::nullopt)},
        {query_fn_name,
         {json1, "$"},
         json_constructor(
             R"({"a":{"b":[{"c":"foo","d":1.23,"f":null}],"e":true}})")},
        {query_fn_name,
         {json1, String("$.a")},
         json_constructor(R"({"b":[{"c":"foo","d":1.23,"f":null}],"e":true})")},
        {query_fn_name,
         {json1, absl::StrCat("$", EscapeKey(sql_standard_mode, "a"))},
         json_constructor(R"({"b":[{"c":"foo","d":1.23,"f":null}],"e":true})")},
        {query_fn_name,
         {json1, String("$.a.b")},
         json_constructor(R"([{"c":"foo","d":1.23,"f":null}])")},
        {query_fn_name,
         {json1, String("$.a.b[0]")},
         json_constructor(R"({"c":"foo","d":1.23,"f":null})")},
        {query_fn_name,
         {json1, String("$.a.b[0].c")},
         json_constructor(R"("foo")")},
        {query_fn_name,
         {json1, String("$.a.b[0].d")},
         json_constructor("1.23")},
        {query_fn_name,
         {json1, String("$.a.b[0].g")},
         json_constructor(absl::nullopt)},
        {query_fn_name,
         {json1, String("$.a.b[1]")},
         json_constructor(absl::nullopt)},
        {query_fn_name,
         {json1, String("$.a.x[0]")},
         json_constructor(absl::nullopt)},
        {query_fn_name,
         {json1, String("$.a.x")},
         json_constructor(absl::nullopt)},
        {query_fn_name, {json1, String("$.a.e")}, json_constructor("true")},
        {query_fn_name,
         {json1, String("abc")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
        {query_fn_name,
         {json1, String("")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
        {query_fn_name,
         {json2, String("$.x")},
         json_constructor("[1,2,3,4,5]")},
        {query_fn_name, {json2, String("$.x[1]")}, json_constructor("2")},
        {query_fn_name,
         {json2, absl::StrCat("$", EscapeKey(sql_standard_mode, "x"), "[1]")},
         json_constructor("2")},
        {query_fn_name,
         {json2, String("$.x[-1]")},
         json_constructor(absl::nullopt)},
        {query_fn_name,
         {json2, absl::StrCat("$.x", EscapeKey(sql_standard_mode, "a"))},
         json_constructor(absl::nullopt)},
        {query_fn_name,
         {json2, absl::StrCat("$.x", EscapeKey(sql_standard_mode, "1"))},
         json_constructor("2")},
        {query_fn_name, {json2, String("$.x[ 1]")}, json_constructor("2")},
        {query_fn_name,
         {json2, String("$.x[10]")},
         json_constructor(absl::nullopt)},
        {query_fn_name,
         {json2, String("$.y.a")},
         json_constructor(absl::nullopt)},
        {query_fn_name,
         {json2, String("$.y[0].a")},
         json_constructor(R"("bar")")},
        {query_fn_name,
         {json2, absl::StrCat("$.y[0]", EscapeKey(sql_standard_mode, "a"))},
         json_constructor(R"("bar")")},
        {query_fn_name,
         {json2, String("$.y[0].b")},
         json_constructor(absl::nullopt)},
        {query_fn_name,
         {json2, String("$.y[1].a")},
         json_constructor(absl::nullopt)},
        {query_fn_name,
         {json2, String("$.y[1].b")},
         json_constructor(R"("baz")")},
        {query_fn_name,
         {json3, String("$")},
         json_constructor(R"({"a.b.c":5})")},
        // Query with dots in the middle of the key
        {query_fn_name,
         {json3, absl::StrCat("$", EscapeKey(sql_standard_mode, "a.b.c"))},
         json_constructor("5")},
        // If $[a.b.c] is a valid path then it should return "5"
        // Else it should return OUT_OF_RANGE current code fails to
        // handle this correctly.
        //
        // {query_fn_name, {json3, String("$[a.b.c]")}, NullString(),
        //  OUT_OF_RANGE},
        {query_fn_name,
         {json4, String("$.longer_field_name")},
         json_constructor("[]")},
        {query_fn_name, {json5, String("$")}, json_constructor("true")},
        {query_fn_name,
         {json7, String("$")},
         json_constructor(R"({"a":{"b":{"c":{"d":3}}}})")},
        {query_fn_name,
         {json8, String("$")},
         json_constructor(R"({"x":[],"y":[1,2,5,3,4]})")},
        {query_fn_name, {json8, String("$.x")}, json_constructor(R"([])")},
        // Deep/wide json.
        {query_fn_name,
         {deep_json, String("$.a.b.c.d.e.f.g.h.i.j")},
         json_constructor(
             R"({"k":{"l":{"m":{"x":"foo","y":10,"z":[1,2,3]}}}})")},
        {query_fn_name,
         {deep_json, String("$.a.b.c.d.e.f.g.h.i.j.k.l.m")},
         json_constructor(R"({"x":"foo","y":10,"z":[1,2,3]})")},
        {query_fn_name,
         {deep_json, String("$.a.b.c.d.e.f.g.h.i.j.k.l.m.z[1]")},
         json_constructor("2")},
        {query_fn_name,
         {deep_json, String("$.a.b.c.d.e.f.g.h.i.j.k.l.m.x")},
         json_constructor(R"("foo")")},
        {query_fn_name, {wide_json, String("$.j")}, json_constructor("[-3,0]")},
        {query_fn_name,
         {array_of_deep_json, String("$.arr[13].index")},
         json_constructor("13")},
        {query_fn_name,
         {array_of_deep_json,
          String("$.arr[13].a.b.c.d.e.f.g.h.i.j.k.l.m.z[1]")},
         json_constructor("2")},
        {query_fn_name,
         {array_of_wide_json, String("$.arr[17].index")},
         json_constructor("17")},
        {query_fn_name,
         {array_of_wide_json, String("$.arr[17].k")},
         json_constructor("0.321")},
        {query_fn_name,
         {array_of_wide_json, String("$.arr[14]")},
         json_constructor(
             R"({"index":14,"a":null,"b":"bar","c":false,"d":[4,5],"e":0.123,)"
             R"("f":"345","g":null,"h":"baz","i":true,"j":[-3,0],"k":0.321,)"
             R"("l":"678"})")},
        // Non-ASCII UTF-8 and special cases.
        {query_fn_name,
         {json_constructor(R"({"Моша_öá5ホバークラフト鰻鰻" : "x"})"),
          String("$")},
         json_constructor(R"({"Моша_öá5ホバークラフト鰻鰻":"x"})")},
        // Unsupported/unimplemented JSONPath features.
        {query_fn_name,
         {json1, String("$.a.*")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
        {query_fn_name,
         {json1, String("$.a.b..c")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
        {query_fn_name,
         {json2, String("$.x[(@.length-1)]")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
        {query_fn_name,
         {json2, String("$.x[-1:]")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
        {query_fn_name,
         {json2, String("$.x[0:4:2]")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
        {query_fn_name,
         {json2, String("$.x[:2]")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
        {query_fn_name,
         {json2, String("$.x[0,1]")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
        {query_fn_name,
         {json2, String("$.y[?(@.a)]")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
        {query_fn_name,
         {json2, String(R"($.y[?(@.a==='bar')])")},
         json_constructor(absl::nullopt),
         OUT_OF_RANGE},
    };
    if (sql_standard_mode) {
      all_tests.push_back({query_fn_name,
                           {json3, String("$['a.b.c']")},
                           json_constructor(absl::nullopt),
                           OUT_OF_RANGE});
    } else {
      all_tests.push_back({query_fn_name,
                           {json3, String("$.\"a.b.c\"")},
                           json_constructor(absl::nullopt),
                           OUT_OF_RANGE});
    }
  }
  return all_tests;
}

const std::vector<FunctionTestCall> GetStringJsonTests(bool sql_standard_mode,
                                                       bool scalar_test_cases) {
  std::string query_fn_name;
  std::string value_fn_name;
  if (sql_standard_mode) {
    query_fn_name = "json_query";
    value_fn_name = "json_value";
  } else {
    query_fn_name = "json_extract";
    value_fn_name = "json_extract_scalar";
  }
  // Wide numbers
  const Value json_with_wide_numbers = String(
      R"({
           "x":11111111111111111111,
           "y":3.14e314,
           "z":123456789012345678901234567890,
           "a":true,
           "s":"foo"
         })");

  std::vector<FunctionTestCall> tests =
      GetJsonTestsCommon(sql_standard_mode, scalar_test_cases,
                         [](absl::optional<absl::string_view> input) {
                           if (input.has_value()) return String(input.value());
                           return NullString();
                         });

  if (scalar_test_cases) {
    tests.push_back({value_fn_name,
                     {json_with_wide_numbers, String("$.x")},
                     String("11111111111111111111")});
    tests.push_back({value_fn_name,
                     {json_with_wide_numbers, String("$.y")},
                     String("3.14e314")});
    tests.push_back({value_fn_name,
                     {json_with_wide_numbers, String("$.z")},
                     String("123456789012345678901234567890")});
    tests.push_back({value_fn_name,
                     {json_with_wide_numbers, String("$.s")},
                     String("foo")});
    // Characters should _not_ be escaped in JSON_EXTRACT_SCALAR.
    tests.push_back(
        {value_fn_name,
         {String(R"({"a": "foo\t\\t\\\t\n\\nbar\ \"baz\\"})"), String("$.a")},
         String("foo\t\\t\\\t\n\\nbar \"baz\\")});
  } else {
    tests.push_back(
        {query_fn_name,
         {String(
              R"({"a":{"b":[{"c" : "foo", "d": 1.23, "f":null }], "e": true}})"),
          String("$.a.b[0].f")},
         NullString()});
    // Malformed JSON.
    tests.push_back(
        {query_fn_name, {String(R"({"a": )"), String("$")}, NullString()});
    // Wide numbers
    tests.push_back(
        {query_fn_name,
         {json_with_wide_numbers, String("$")},
         String(R"({"x":11111111111111111111,"y":3.14e314,)"
                R"("z":123456789012345678901234567890,"a":true,"s":"foo"})")});
    // Characters should be escaped in JSON_EXTRACT.
    tests.push_back(
        {query_fn_name,
         {String(R"({"a": "foo\t\\t\\\t\n\\nbar\ \"baz\\"})"), String("$.a")},
         String(R"("foo\t\\t\\\t\n\\nbar \"baz\\")")});
  }
  return tests;
}

const std::vector<FunctionTestCall> GetNativeJsonTests(
    bool sql_standard_mode, bool scalar_test_cases) {
  std::string query_fn_name;
  std::string value_fn_name;
  if (sql_standard_mode) {
    query_fn_name = "json_query";
    value_fn_name = "json_value";
  } else {
    query_fn_name = "json_extract";
    value_fn_name = "json_extract_scalar";
  }

  // Wide numbers
  const Value json_with_wide_numbers =
      Json(JSONValue::ParseJSONString(R"({"x":11111111111111111111, )"
                                      R"("z":123456789012345678901234567890, )"
                                      R"("a":true, "s":"foo"})")
               .ValueOrDie());

  std::vector<FunctionTestCall> tests = GetJsonTestsCommon(
      sql_standard_mode, scalar_test_cases,
      [](absl::optional<absl::string_view> input) {
        if (input.has_value())
          return Json(JSONValue::ParseJSONString(input.value()).ValueOrDie());
        return NullJson();
      });

  if (scalar_test_cases) {
    tests.push_back({value_fn_name,
                     {json_with_wide_numbers, String("$.x")},
                     String("11111111111111111111")});
    tests.push_back({value_fn_name,
                     {json_with_wide_numbers, String("$.z")},
                     String("1.2345678901234568e+29")});
    tests.push_back({value_fn_name,
                     {json_with_wide_numbers, String("$.s")},
                     String("foo")});
    // Characters should _not_ be escaped in JSON_EXTRACT_SCALAR.
    tests.push_back({value_fn_name,
                     {Json(JSONValue::ParseJSONString(
                               R"({"a": "foo\t\\t\\\t\n\\nbar \"baz\\"})")
                               .ValueOrDie()),
                      String("$.a")},
                     String("foo\t\\t\\\t\n\\nbar \"baz\\")});
  } else {
    JSONValue null_json_value;
    // This case differs from the string version of STRING_EXTRACT/QUERY. The
    // result is a JSON 'null' which is different from a null JSON.
    tests.push_back(
        {query_fn_name,
         {Json(JSONValue::ParseJSONString(
                   R"({"a":{"b":[{"c" : "foo", "d": 1.23, "f":null }], )"
                   R"("e": true}})")
                   .ValueOrDie()),
          String("$.a.b[0].f")},
         Json(std::move(null_json_value))});
    // Malformed JSON.
    tests.push_back({query_fn_name,
                     {Value::UnvalidatedJsonString(R"({"a": )"), String("$")},
                     NullJson(),
                     OUT_OF_RANGE});
    tests.push_back(
        {query_fn_name,
         {json_with_wide_numbers, String("$")},
         Json(JSONValue::ParseJSONString(
                  R"({"x":11111111111111111111,)"
                  R"("z":123456789012345678901234567890,"a":true,"s":"foo"})")
                  .ValueOrDie())});
    // Characters should be escaped in JSON_EXTRACT.
    tests.push_back(
        {query_fn_name,
         {Json(JSONValue::ParseJSONString(
                   R"({"a": "foo\t\\t\\\t\n\\nbar \"baz\\"})")
                   .ValueOrDie()),
          String("$.a")},
         Json(JSONValue::ParseJSONString(R"("foo\t\\t\\\t\n\\nbar \"baz\\")")
                  .ValueOrDie())});
  }

  return tests;
}

}  // namespace

std::vector<FunctionTestCall> GetFunctionTestsStringJsonQuery() {
  return GetStringJsonTests(/*sql_standard_mode=*/true,
                            /*scalar_test_cases=*/false);
}

std::vector<FunctionTestCall> GetFunctionTestsStringJsonExtract() {
  return GetStringJsonTests(/*sql_standard_mode=*/false,
                            /*scalar_test_cases=*/false);
}

std::vector<FunctionTestCall> GetFunctionTestsStringJsonValue() {
  return GetStringJsonTests(/*sql_standard_mode=*/true,
                            /*scalar_test_cases=*/true);
}

std::vector<FunctionTestCall> GetFunctionTestsStringJsonExtractScalar() {
  return GetStringJsonTests(/*sql_standard_mode=*/false,
                            /*scalar_test_cases=*/true);
}

std::vector<FunctionTestCall> GetFunctionTestsNativeJsonQuery() {
  return GetNativeJsonTests(/*sql_standard_mode=*/true,
                            /*scalar_test_cases=*/false);
}

std::vector<FunctionTestCall> GetFunctionTestsNativeJsonExtract() {
  return GetNativeJsonTests(/*sql_standard_mode=*/false,
                            /*scalar_test_cases=*/false);
}

std::vector<FunctionTestCall> GetFunctionTestsNativeJsonValue() {
  return GetNativeJsonTests(/*sql_standard_mode=*/true,
                            /*scalar_test_cases=*/true);
}

std::vector<FunctionTestCall> GetFunctionTestsNativeJsonExtractScalar() {
  return GetNativeJsonTests(/*sql_standard_mode=*/false,
                            /*scalar_test_cases=*/true);
}

std::vector<QueryParamsWithResult> GetFunctionTestsJsonIsNull() {
  std::vector<QueryParamsWithResult> v = {
      {{NullJson()}, True()},
      {{Json(JSONValue(1.1))}, False()},
      {{Json(JSONValue(std::string{"null"}))}, False()},
      {{Value::Null(types::JsonArrayType())}, True()},
  };
  return v;
}

}  // namespace zetasql
